import matplotlib.pyplot as plt
import pandas as pd
from datetime import datetime, timedelta
from database import Database
from categories import EXPENSE_CATEGORIES, INCOME_CATEGORIES
import tkinter as tk
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
import matplotlib
matplotlib.use('TkAgg')  # 确保使用TkAgg后端

class FinancialAnalysis:
    def __init__(self, db: Database):
        self.db = db
    
    def get_monthly_analysis(self, year: int, month: int):
        """获取月度消费分析 - 优化版本"""
        # 使用datetime处理日期避免边界问题
        start_date = datetime(year, month, 1)
        if month == 12:
            end_date = datetime(year + 1, 1, 1)
        else:
            end_date = datetime(year, month + 1, 1)
        
        # 参数化查询防止SQL注入
        transactions = self.db.get_transactions(
            start_date.strftime("%Y-%m-%d"),
            end_date.strftime("%Y-%m-%d")
        )
        
        total_income = sum(t['amount'] for t in transactions if t['type'] == 'income')
        total_expense = sum(t['amount'] for t in transactions if t['type'] == 'expense')
        balance = total_income - total_expense
        
        # 预处理交易数据，避免多次遍历
        expense_by_category = {
            'main': {},
            'middle': {},
            'sub': {}
        }
        
        for t in transactions:
            if t['type'] != 'expense':
                continue
                
            main_cat = t['main_category']
            sub_cat_full = t['sub_category']
            
            # 解析子分类结构
            parts = sub_cat_full.split(" - ")
            middle_cat = parts[0] if len(parts) >= 1 else ""
            sub_cat = parts[1] if len(parts) >= 2 else ""
            
            # 更新各类别统计数据
            expense_by_category['main'][main_cat] = \
                expense_by_category['main'].get(main_cat, 0) + t['amount']
                
            expense_by_category['middle'][middle_cat] = \
                expense_by_category['middle'].get(middle_cat, 0) + t['amount']
                
            expense_by_category['sub'][sub_cat] = \
                expense_by_category['sub'].get(sub_cat, 0) + t['amount']
        
        # 过滤掉金额为0的项目
        main_category_expense = {k: v for k, v in expense_by_category['main'].items() if v > 0}
        middle_category_expense = {k: v for k, v in expense_by_category['middle'].items() if v > 0}
        sub_category_expense = {k: v for k, v in expense_by_category['sub'].items() if v > 0}
        
        savings_rate = 0
        if total_income > 0:
            try:
                savings_rate = (balance / total_income) * 100
            except ZeroDivisionError:
                pass  # 已经检查过，理论上不会触发
        
        return {
            'total_income': total_income,
            'total_expense': total_expense,
            'balance': balance,
            'main_category_expense': main_category_expense,
            'middle_category_expense': middle_category_expense,
            'sub_category_expense': sub_category_expense,
            'savings_rate': savings_rate
        }

    def _create_pie_chart(self, parent_frame, category_data, title):
        """辅助函数：绘制饼图"""
        if not category_data:
            label = tk.Label(parent_frame, text=title.replace("分布", "暂无消费数据"), font=("Arial", 12))
            label.pack(pady=20)
            return None
            
        # 如果分类太多，只显示前8个，其他的合并为"其他"
        if len(category_data) > 8:
            sorted_categories = sorted(category_data.items(), key=lambda x: x[1], reverse=True)
            top_categories = dict(sorted_categories[:7])
            other_total = sum(item[1] for item in sorted_categories[7:])
            if other_total > 0:
                top_categories['其他'] = other_total
            category_data = top_categories
            
        fig, ax = plt.subplots(figsize=(10, 8))
        
        # 设置中文字体
        plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'DejaVu Sans']
        plt.rcParams['axes.unicode_minus'] = False
        
        # 使用更鲜明的颜色
        colors = plt.cm.tab20(range(len(category_data)))
        
        # 绘制饼图，优化标签显示
        wedges, texts, autotexts = ax.pie(
            category_data.values(), 
            labels=category_data.keys(),
            autopct=lambda pct: f'{pct:.1f}%' if pct >= 5 else '',  # 只显示大于5%的百分比
            colors=colors,
            startangle=90,
            textprops={'fontsize': 9},
            pctdistance=0.85,  # 百分比文字距离圆心的距离
            labeldistance=1.05  # 标签距离圆心的距离
        )
        
        # 美化文本
        for autotext in autotexts:
            autotext.set_color('black')
            autotext.set_fontweight('bold')
            autotext.set_fontsize(10)
        
        # 设置标题
        ax.set_title(title, fontsize=14, fontweight='bold', pad=20)
        
        # 添加图例（对于小百分比的部分）
        if any(pct < 5 for pct in [100 * value / sum(category_data.values()) for value in category_data.values()]):
            # 创建图例，显示所有类别的金额
            legend_labels = [f'{label}: {amount:.2f}元' for label, amount in category_data.items()]
            ax.legend(wedges, legend_labels, title="分类详情", loc="center left", bbox_to_anchor=(1, 0, 0.5, 1))
        
        # 调整布局，为图例留出空间
        plt.tight_layout()
        
        # 嵌入到Tkinter窗口
        canvas = FigureCanvasTkAgg(fig, parent_frame)
        canvas.draw()
        canvas.get_tk_widget().pack(fill=tk.BOTH, expand=True)
        
        return canvas

    def plot_category_distribution(self, parent_frame, year: int, month: int, level: str = "main"):
        """绘制消费分类饼图 - 优化版本"""
        try:
            analysis = self.get_monthly_analysis(year, month)
            
            if level == "main":
                category_data = analysis['main_category_expense']
                title = f'{year}年{month}月消费大类分布'
            elif level == "middle":
                category_data = analysis['middle_category_expense']
                title = f'{year}年{month}月消费中类分布'
            else:  # sub
                category_data = analysis['sub_category_expense']
                title = f'{year}年{month}月消费小类分布'

            return self._create_pie_chart(parent_frame, category_data, title)
            
        except Exception as e:
            # 错误处理
            error_label = tk.Label(parent_frame, text=f"生成图表时出错: {str(e)}", font=("Arial", 12), foreground="red")
            error_label.pack(pady=20)
            return None
        
    def plot_monthly_trend(self, parent_frame, months: int = 6):
        """绘制月度收支趋势图 - 优化版本"""
        try:
            monthly_data = self.db.get_summary_by_period('month')
            
            if not monthly_data:
                label = tk.Label(parent_frame, text="暂无数据，无法生成趋势图", font=("Arial", 12))
                label.pack(pady=20)
                return None
            
            # 限制显示的月份数量
            monthly_data = monthly_data[:months]
            
            if len(monthly_data) < 2:
                label = tk.Label(parent_frame, text="数据不足，无法生成趋势图", font=("Arial", 12))
                label.pack(pady=20)
                return None
            
            periods = [data['period'] for data in monthly_data][::-1]
            incomes = [data['total_income'] for data in monthly_data][::-1]
            expenses = [data['total_expense'] for data in monthly_data][::-1]
            
            fig, ax = plt.subplots(figsize=(10, 6))
            
            # 设置中文字体
            plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei']
            plt.rcParams['axes.unicode_minus'] = False
            
            x = range(len(periods))
            width = 0.35
            
            ax.bar([i - width/2 for i in x], incomes, width, label='收入', alpha=0.8, color='lightgreen')
            ax.bar([i + width/2 for i in x], expenses, width, label='支出', alpha=0.8, color='lightcoral')
            
            ax.set_xlabel('月份')
            ax.set_ylabel('金额 (元)')
            ax.set_title('月度收支趋势')
            ax.set_xticks(x)
            ax.set_xticklabels(periods, rotation=45)
            ax.legend()
            
            # 在柱状图上显示数值
            for i, v in enumerate(incomes):
                if v > 0:
                    ax.text(i - width/2, v + max(incomes)/50, f'{v:.0f}', ha='center', va='bottom', fontsize=8)
            for i, v in enumerate(expenses):
                if v > 0:
                    ax.text(i + width/2, v + max(expenses)/50, f'{v:.0f}', ha='center', va='bottom', fontsize=8)
            
            plt.tight_layout()
            
            canvas = FigureCanvasTkAgg(fig, parent_frame)
            canvas.draw()
            canvas.get_tk_widget().pack(fill=tk.BOTH, expand=True)
            
            return canvas
            
        except Exception as e:
            error_label = tk.Label(parent_frame, text=f"生成趋势图时出错: {str(e)}", font=("Arial", 12), foreground="red")
            error_label.pack(pady=20)
            return None        